%% ============================================================
%  Train on Z_all.csv (trained dataset only)
%  Test on untrained_USV_Z.csv and untrained_nonUSV_Z.csv
%  Pipeline:
%    1) Load training data (Z_all) and labels (non-USV then USV)
%    2) Fit PCA on TRAIN only (PC1–PC10)
%    3) Train linear SVM on TRAIN PCs (tune C by 5-fold CV on TRAIN only)
%    4) Transform TEST using TRAIN PCA (no leakage)
%    5) Evaluate on TEST: confusion matrix, accuracy, FPR/FNR, ROC/AUC
%  Output figures: ROC + confusion matrix
%% ============================================================

clear; clc; close all;

%% --------- User settings ----------
train_file = 'Z_all.csv';
test_nonusv_file = 'untrained_nonUSV_Z.csv';
test_usv_file    = 'untrained_USV_Z.csv';

% IMPORTANT: these numbers must match how Z_all was constructed:
n_train_nonUSV = 277;   % first N rows in training are non-USV
n_train_USV    = 171;   % last  N rows in training are USV

nPC = 10;
kfold = 5;
C_grid = [0.01 0.1 1 10 100];

% If you want deterministic t-SNE/UMAP later, set rng here; not needed for PCA/SVM.
rng(1);

%% --------- Load TRAIN (Z_all.csv) ----------
% Z_all structure: first row = time (with first entry 'ID'), first column = ID, data start at row2/col2
opts = detectImportOptions(train_file);
opts.DataLines = 2;                 % skip header row (timestamps)
raw_train = readmatrix(train_file, opts);

% Remove ID column
X_train = raw_train(:, 2:end);

% Labels assume Z_all rows are ordered: non-USV first, USV last
if size(X_train,1) ~= (n_train_nonUSV + n_train_USV)
    error('Training rows (%d) do not match n_train_nonUSV+n_train_USV (%d). Update counts.', ...
        size(X_train,1), n_train_nonUSV + n_train_USV);
end
y_train = [zeros(n_train_nonUSV,1); ones(n_train_USV,1)];

% (Optional) sanity check for time bins count
fprintf('TRAIN: %d events, %d time bins\n', size(X_train,1), size(X_train,2));

%% --------- Fit PCA on TRAIN only ----------
% MATLAB pca centers columns by default and returns the training mean in mu
[coeff, score_train, latent, ~, explained, mu] = pca(X_train);

PC_train = score_train(:, 1:nPC);

fprintf('PCA TRAIN: cumulative variance explained by first %d PCs = %.2f%%\n', ...
    nPC, sum(explained(1:nPC)));

%% --------- Tune linear SVM C (TRAIN only) ----------
cvp = cvpartition(y_train, 'KFold', kfold);

bestC = C_grid(1);
bestLoss = inf;

for C = C_grid
    svm_model = fitcsvm(PC_train, y_train, ...
        'KernelFunction','linear', ...
        'BoxConstraint', C, ...
        'Standardize', false, ...   % data already Z-scored in your naming; avoid double-standardization
        'ClassNames', [0 1]);

    svm_cv = crossval(svm_model, 'CVPartition', cvp);
    loss = kfoldLoss(svm_cv);  % classification error

    if loss < bestLoss
        bestLoss = loss;
        bestC = C;
    end
end

fprintf('SVM TRAIN tuning: best C = %.3g (mean CV error = %.4f; CV accuracy = %.2f%%)\n', ...
    bestC, bestLoss, 100*(1-bestLoss));

%% --------- Train final SVM on full TRAIN ----------
svm_final = fitcsvm(PC_train, y_train, ...
    'KernelFunction','linear', ...
    'BoxConstraint', bestC, ...
    'Standardize', false, ...
    'ClassNames', [0 1]);

%% --------- Load TEST (untrained) ----------
% untrained files structure: first row = time, rows 2..end = events
A = readmatrix(test_nonusv_file);
t_test = A(1, :);
X_test_non = A(2:end, :);

B = readmatrix(test_usv_file);
t_test2 = B(1, :);
X_test_usv = B(2:end, :);

% Basic time consistency checks
if numel(t_test) ~= numel(t_test2) || any(abs(t_test - t_test2) > 1e-12)
    error('TEST time vectors differ between non-USV and USV test files.');
end

% Ensure TEST time bins match TRAIN time bins count
if size(X_test_non,2) ~= size(X_train,2) || size(X_test_usv,2) ~= size(X_train,2)
    error('Time bin count mismatch: TRAIN=%d, TEST_non=%d, TEST_usv=%d.', ...
        size(X_train,2), size(X_test_non,2), size(X_test_usv,2));
end

X_test = [X_test_non; X_test_usv];
y_test = [zeros(size(X_test_non,1),1); ones(size(X_test_usv,1),1)];

fprintf('TEST: %d events (non-USV=%d, USV=%d), %d time bins\n', ...
    size(X_test,1), size(X_test_non,1), size(X_test_usv,1), size(X_test,2));

%% --------- Transform TEST using TRAIN PCA (no leakage) ----------
PC_test = (X_test - mu) * coeff(:, 1:nPC);

%% --------- Predict on TEST ----------
[y_pred, score] = predict(svm_final, PC_test);
% score(:,2) corresponds to class 1 when ClassNames=[0 1]

%% --------- Confusion matrix + metrics ----------
cm = confusionmat(y_test, y_pred);  % rows=true, cols=pred
TN = cm(1,1); FP = cm(1,2);
FN = cm(2,1); TP = cm(2,2);

acc = (TP + TN) / sum(cm(:));
fpr = FP / (FP + TN);
fnr = FN / (FN + TP);
precision = TP / max(TP + FP, 1);
recall    = TP / max(TP + FN, 1);

fprintf('\n=== HELD-OUT TEST PERFORMANCE (SVM, linear) ===\n');
fprintf('Accuracy   = %.3f\n', acc);
fprintf('Precision  = %.3f\n', precision);
fprintf('Recall(TPR)= %.3f\n', recall);
fprintf('FPR        = %.3f\n', fpr);
fprintf('FNR        = %.3f\n', fnr);
fprintf('TP=%d  FP=%d  TN=%d  FN=%d\n', TP, FP, TN, FN);

%% --------- ROC/AUC on TEST ----------
[Xroc, Yroc, ~, auc] = perfcurve(y_test, score(:,2), 1);

figure;
plot(Xroc, Yroc, 'LineWidth', 2); grid on;
xlabel('False Positive Rate');
ylabel('True Positive Rate');
title(sprintf('Held-out TEST ROC (Linear SVM, PC1-%d), AUC = %.2f', nPC, auc));
saveas(gcf, sprintf('Supp_SVM_Heldout_ROC_PC1to%d.png', nPC));

%% --------- Confusion matrix figure ----------
figure;
imagesc(cm);
colormap(parula);
colorbar;

xticks([1 2]); yticks([1 2]);
xticklabels({'Pred non-USV','Pred USV'});
yticklabels({'True non-USV','True USV'});
axis image;

title(sprintf('Held-out TEST Confusion Matrix (Acc=%.2f, FPR=%.2f, FNR=%.2f)', acc, fpr, fnr));

% Overlay numbers
text(1,1,num2str(TN),'HorizontalAlignment','center','Color','w','FontSize',14,'FontWeight','bold');
text(2,1,num2str(FP),'HorizontalAlignment','center','Color','w','FontSize',14,'FontWeight','bold');
text(1,2,num2str(FN),'HorizontalAlignment','center','Color','w','FontSize',14,'FontWeight','bold');
text(2,2,num2str(TP),'HorizontalAlignment','center','Color','w','FontSize',14,'FontWeight','bold');

saveas(gcf, sprintf('Supp_SVM_Heldout_ConfusionMatrix_PC1to%d.png', nPC));

%% --------- Save a small summary text file ----------
summary_file = sprintf('Supp_SVM_Heldout_Summary_PC1to%d.txt', nPC);
fid = fopen(summary_file, 'w');
fprintf(fid, 'Linear SVM decoding on held-out TEST set\n');
fprintf(fid, 'Training file: %s\n', train_file);
fprintf(fid, 'Test files: %s | %s\n', test_nonusv_file, test_usv_file);
fprintf(fid, 'PC features: PC1-%d (PCA fit on TRAIN only)\n', nPC);
fprintf(fid, 'Best C (5-fold CV on TRAIN): %.3g\n', bestC);
fprintf(fid, 'Held-out TEST AUC: %.4f\n', auc);
fprintf(fid, 'Accuracy: %.4f\n', acc);
fprintf(fid, 'Precision: %.4f\n', precision);
fprintf(fid, 'Recall/TPR: %.4f\n', recall);
fprintf(fid, 'FPR: %.4f\n', fpr);
fprintf(fid, 'FNR: %.4f\n', fnr);
fprintf(fid, 'TP=%d  FP=%d  TN=%d  FN=%d\n', TP, FP, TN, FN);
fclose(fid);

fprintf('\nSaved:\n  %s\n  %s\n  %s\n', ...
    sprintf('Supp_SVM_Heldout_ROC_PC1to%d.png', nPC), ...
    sprintf('Supp_SVM_Heldout_ConfusionMatrix_PC1to%d.png', nPC), ...
    summary_file);
